from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name="my_reduce",
    ext_modules=[
        CppExtension(
            "my_reduce",
            ["my_reduce.cpp"],  # 对应 PYBIND11_MODULE 名
            extra_compile_args=["-O3"],
        )  # 可加 -march=native / CUDA 标志
    ],
    cmdclass={"build_ext": BuildExtension},
)
